library(Signac)
source('00_util_scripts/mod_seurat.R')
source('00_util_scripts/mod_bplot.R')
library(GenomicRanges)
library(future)
options("restore_Seurat_show" = TRUE)

plan('multisession', workers = 12)

options(future.globals.maxSize = 50000 * 1024^2) # for 50 Gb RAM

proj.nm <- 'mission/FPP/ober23/'

sc.path <- list.files('~/append-ssd/macs3/', pattern = 'singlecell',
           recursive = T, full.names = T) |>
  str_subset('outs') |>
  str_subset('SD2_macs3', negate = T)

lib.name <- sc.path |>
  str_extract('C_\\w{3}|AA\\d')

lib.name

frg.path <- list.files('~/append-ssd/macs3/', pattern = 'fragments.tsv.gz$',
                      full.names = T) |>
  str_subset('POOL', negate = T)

frg.path

tibble(library_id = lib.name,
       fragments = frg.path,
       cells = sc.path) |>
  write_csv('~/append-ssd/macs3/aggr10.csv')

# merge peaks -----------
# read in peak sets
peak.path <- list.files('~/append-ssd/macs3/', pattern = 'final_peaks',
                      recursive = T, full.names = T) |>
  str_subset('POOL|macs3\\.', negate = T)

peak.path

peaks.pb1 <- read_delim(
  '~/append-ssd/macs3/GSE212448_C_PB1/final_peaks.bed',
  col_names = c("chr", "start", "end")
)

peaks.pb1

gr.cat <- peak.path |>
read_delim(col_names = c("chr", "start", "end")) |>
  makeGRangesFromDataFrame()

# Create a unified set of peaks to quantify in each dataset
combined.peaks <- GenomicRanges::reduce(gr.cat)

# Filter out bad peaks based on length
peakwidths <- width(combined.peaks)

combined.peaks <- combined.peaks[peakwidths  < 10000 & peakwidths > 20]
combined.peaks

# load metadata --------------
# perform an initial filtering of low count cells
sc.meta <- sc.path |>
  map(\(x)read_delim(x) |> filter(passed_filters > 1000))

sc.meta

# keep only supp-provided barcodes
ober.meta <- read_tsv('mission/FPP/ober23/ober23.atac.meta.tsv',
                      skip = 1)

ober.meta |>
  select(...1) |>
  filter(!str_detect(...1, 'POOL'))

# create fragment objects ---------
fobj.list <- frg.path |>
  map2(sc.meta,
       \(x,y)CreateFragmentObject(x, y$barcode))

# quantify peak as fmtx ----------
# 40 multisession = 7.8min
# 16 multisession = 5.1min
# 12 ms = 4.4m
# 8 ms = 4.5m
# 4 multisession = 5.9min
# sequential = 20min
# multi + 10 sample = 40min
fmtx.list <- fobj.list |>
  map2(sc.meta, \(x,y)FeatureMatrix(x, combined.peaks, y$barcode),
       .progress = T)

fobj.list[[1]]

system.time(fmtx1 <- FeatureMatrix(fobj.list[[1]],
                       combined.peaks,
                       sc.meta[[1]]$barcode))

aobj1 <- fmtx1 |>
  CreateChromatinAssay(fragments = fobj1) |>
  CreateSeuratObject(assay = 'ATAC',
                     meta.data = column_to_rownames(sc.meta[[1]], 'barcode'))

aobj.list <- fmtx.list |>
  map2(fobj.list, \(x,y)CreateChromatinAssay(x,fragments = y)) |>
  map2(sc.meta, \(x,y)CreateSeuratObject(x, assay = 'ATAC',
                                         meta.data = column_to_rownames(y, 'barcode')),
       .progress = T)

# add information to identify dataset of origin
aobj.list <- aobj.list |> map2(lib.name, \(x,y)mutate(x, dataset = y))

aobj.list[[10]]@meta.data |> as_tibble(rownames = '.cell')

# merge all datasets, adding a cell ID to make sure cell names are unique
combined <- merge(
  x = aobj.list[[1]],
  y = aobj.list[-1],
  add.cell.ids = lib.name
)
# 371k cells
combined[["ATAC"]]

combined@meta.data |> as_tibble(rownames = '.cell')

# 40k cells post-filtering
aobj.filtered <- combined |>
  mutate(barcode = str_remove(.cell, '.+_'),
         supp.barcode = str_c(dataset, '#', barcode)) |>
  filter(supp.barcode %in% ober.meta$...1)

aobj.filtered

aobj.filtered <- aobj.filtered |>
  mutate(group = if_else(str_detect(Sample, 'AA'), 'AA', 'HC')) |>
  mutate(subgroup = str_c(NamedClust, "-", group))

aobj.kc <- aobj.kc |>
  mutate(subgroup = str_c(NamedClust, "-", group))

# big rds -----------
aobj.filtered |> write_rds('mission/FPP/ober23/ober23.s10.rds')

aobj.filtered <- read_rds('mission/FPP/ober23/ober23.s10.rds')

aobj.hc <- aobj.filtered |>
  filter(str_detect(dataset, '^C'))

aobj.hc |> write_rds('mission/FPP/ober23/ober23.hc.rds')

aobj.psor <- read_rds('mission/FPP/psoriasis/sun24.ctrl.signac.rds')

aobj.psor |> CoveragePlot('TRPV3', group.by = 'manual.main')

Annotation(aobj.filtered) <- Annotation(aobj.psor)

rm(aobj.psor)

# LSI (TFIDF Normalization + SVD dim reduc) -------
aobj.filtered <- aobj.filtered |>
  RunTFIDF() |>
  FindTopFeatures() |>
  RunSVD()

#  component 1 will strongly correlate with seq depth
# So we will exclude it from downstream analysis
DepthCor(aobj.filtered)

# SLM (smart local moving) clustering
aobj.filtered <- aobj.filtered |>
  RunUMAP(reduction = 'lsi', dims = 2:30) |>
  FindNeighbors(reduction = 'lsi', dims = 2:30) |>
  FindClusters(algorithm = 3)

aobj.filtered <- aobj.filtered |>
  left_join(dplyr::rename(ober.meta, 'supp.barcode' = '...1'))

aobj.filtered <- aobj.filtered |>
  dplyr::select(-any_of(colnames(ober.meta)))

aobj.filtered |>
  DimPlot(label = TRUE, label.box = T, cols = DiscretePalette(36),
          group.by = 'BroadClust')

aobj.kc <- aobj.filtered |>
  filter(BroadClust == 'Kc', group == 'HC')

aobj.kc |>
  DimPlot(label = TRUE, label.box = T, cols = DiscretePalette(36),
          group.by = 'FineClust')

# estimate gene activity ----------------
# may cost ~7min for all genes in HC KC with 5e3 processes_n
# cost 13~14min for all genes in HC with 5e3 processes_n (30G RAM usage)
# can specify interested genes
# (but proper normalization need whole-transcriptome evaluation)
system.time(gna.kc <- GeneActivity(aobj.kc, process_n = 5e3))

aobj.kc[['RNA']] <- gna.kc |>
  CreateAssayObject()

aobj.kc <- aobj.kc |> NormalizeData(
  assay = 'RNA',
  normalization.method = 'LogNormalize',
  scale.factor = median(aobj.kc$nCount_RNA)
)

## ADRB2 activity ------------
system.time(gna <- GeneActivity(aobj.filtered, process_n = 5e3))

aobj.filtered[['RNA']] <- gna |>
  CreateAssayObject()

gna_median <- median(aobj.filtered$nCount_RNA)

aobj.filtered <- aobj.filtered |> NormalizeData(
  assay = 'RNA',
  normalization.method = 'LogNormalize',
  scale.factor = gna_median
)

aobj.filtered <- aobj.filtered |>
  mutate(type_main = case_match(BroadClust,
                                'Kc' ~ 'Keratinocytes',
                                'Fb' ~ 'Fibroblasts',
                                'Ve' ~ 'Endothelial cells',
                                'Tc' ~ 'T cells',
                                'My' ~ 'Myeloid cells',
                                'Mu' ~ 'Muscle cells',
                                'Me' ~ 'Melanocytes',
                                'Le' ~ 'Lymphatics',
                                'Bc' ~ 'B cells'))

adr.hs.genes <- gna |> rownames() |> str_subset('ADR[A-C]')

aobj.hc |>
  DotPlot(adr.hs.genes, assay = 'RNA', cols = 'RdYlBu',
                         group.by = 'type_main') +
  RotatedAxis() +
  labs(x = 'Gene', y = 'Cell type',
       title = 'Gene activity in human healthy skin scATAC-seq data')

## TRPV3 activity -----------
aobj.kc.v3 |>
  DotPlot(c('TRPV3'), group.by = 'FineClust',
          cols = 'RdYlBu', assay = 'RNA')

Idents(aobj.kc) <- 'FineClust'

aobj.kc |>
  FindAllMarkers('TRPV3', only.pos = T, assay = 'RNA')

aobj.kc |>
  DotPlot(c('HMGCS1'), group.by = 'FineClust',
          cols = 'RdYlBu', assay = 'RNA')

aobj.kc |>
  FindAllMarkers('HMGCS1', only.pos = T, assay = 'RNA')

## MVA pathway activity ------------
kegg_mva <-
  c('Hmgcs1','Hmgcr','Mvk','Pmvk','Fdps') |> str_to_upper()

aobj.kc |>
  DotPlot(kegg_mva, group.by = 'FineClust',
          cols = 'RdYlBu',assay = 'RNA') +
  RotatedAxis()

Idents(aobj.kc) <- 'FineClust'

aobj.kc |>
  FindAllMarkers(assay = 'RNA', features = 'TRPV3', only.pos = T)

gna.v3 <- aobj.kc |>
  GeneActivity(features = 'TRPV3')

aobj.filtered <- gna.v3 |>
  as_tibble(rownames = '.cell') |>
  dplyr::rename(TRPV3.count = value) |>
  left_join(x = aobj.filtered, y = _)

aobj.filtered |>
  summarise(v3.mean = mean(TRPV3.count), .by = BroadClust)

# motif analysis -----------
library(TFBSTools)

pfm <- getMatrixSet(
  x = '~/append-ssd/work/JASPAR2024.sqlite',
  opts = list(collection = "CORE", tax_group = 'vertebrates',
              all_versions = FALSE)
)

aobj.filtered <- aobj.filtered |>
  AddMotifs(pfm = pfm,
            genome = BSgenome.Hsapiens.UCSC.hg38::BSgenome.Hsapiens.UCSC.hg38)

Idents(aobj.filtered) <- 'NamedClust'

v3h.dep <- aobj.filtered |> FindMarkers(
  ident.1 = c('aKc1','aKc2'),
  group.by = 'NamedClust',
  min.pct = 0.05,
  assay = 'ATAC'
)

# get top differentially accessible peaks
top.da.peak <- da_peaks |>
  as_tibble(rownames = 'peak') |>
  dplyr::filter(p_val < .005, pct.1 > .2) |>
  pull(peak)

## check motif in peak ----------
motif.mtx <- aobj.kc.v3 |> GetMotifData()
motif.mtx |> glimpse()
motif.mtx[1:5,1:5]
srebf.mtx <- motif.mtx[,motif.srebf$name]
srebf.mtx.rs <- srebf.mtx |> rowSums2()

srebf.mtx <- srebf.mtx[srebf.mtx.rs > 0, ] |>
  as_tibble(rownames = 'region') |>
  separate_wider_delim(region, '-', names = c('chr','start', 'end'))

srebf.mtx

srebf.mtx |> filter(chr == 'chr5', start > 4329e4, start < 4333e4)

# chromvar --------
# cost 10+ min
aobj.filtered <- aobj.filtered |>
  RunChromVAR(genome = BSgenome.Hsapiens.UCSC.hg38::BSgenome.Hsapiens.UCSC.hg38)

aobj.filtered |>
  save(file = 'mission/FPP/ober23/ober23.chrvr.rda')

load('mission/FPP/ober23/ober23.chrvr.rda')

motif.meta <- aobj.kc |>
  GetMotifData(slot = 'motif.names') |>
  as_tibble() |>
  pivot_longer(everything())

motif.srebf <- motif.meta |>
  dplyr::filter(str_detect(value, 'SREBF'))

motif.srebf2 <- motif.meta |>
  dplyr::filter(str_detect(value, 'SREBF2'))

motif.ctcf <- motif.meta |>
  dplyr::filter(str_detect(value, 'CTCF($|\\.)'))

aobj.kc |>
  DotPlot('TRPV3', assay = 'RNA', cols = 'RdYlBu',
          group.by = 'NamedClust')

aobj.kc |>
  DotPlot(assay = 'chromvar',
          features = motif.srebf$name,
          group.by = 'FineClust',
          cols = 'RdYlBu') +
  RotatedAxis() +
  labs(title = 'SREBF2 binding motif openness', x = 'Motif name',
       y = 'Cell type')

# volcano ------------
aobj.v3hl.fine <- aobj.filtered |>
  filter(str_detect(FineClust, 'aKc1$|aKc2|aKc3')) |>
  mutate(v3.type = ifelse(FineClust == 'aKc1', 'V3-lo KC', 'V3-hi KC'))

aobj.v3hl.named <- aobj.filtered |>
  filter(str_detect(NamedClust, 'aKc1$|aKc2|aKc3')) |>
  mutate(v3.type = ifelse(NamedClust == 'aKc1', 'V3-lo KC', 'V3-hi KC'))

aobj.v3hl.fine$FineClust |> table()

aobj.v3hl.named$NamedClust |> table()

v3h.chrvr <- aobj.kc |>
  FindMarkers(
    ident.1 = c('aKc7','aKc9'),
    ident.2 = NULL,
    group.by = 'FineClust',
    mean.fxn = rowMeans,
    fc.name = "avg_diff",
    assay = 'chromvar'
  )

v3h.chrvr |>
  as_tibble(rownames = 'gene')
  
v3h.chrvr |>
  as_tibble(rownames = 'gene') |>
  mutate(avg_log2FC = avg_diff) |>
  plot_bill_volc(group1 = 'Spinous KC', group2 = 'Basal KC',
                 highlights = motif.srebf$name, force = T)

# footprint -------------
# do not try flanks too short!
# take 17m if in.peaks = T
system.time(
  aobj.kc <- aobj.kc |>
    Footprint(motif.name = c("SREBF1", "CTCF"),
              in.peaks = T, upstream = 250, downstream = 250,
              genome = BSgenome.Hsapiens.UCSC.hg38::BSgenome.Hsapiens.UCSC.hg38)
)

g2 <- aobj.kc |>
  PlotFootprint(c("SREBF1", "CEBPA"), group.by = 'FineClust') +
  patchwork::plot_layout(ncol = 1)

g2

# coverage plot =========
aobj.kc.v3 <- aobj.kc |>
  filter(FineClust %in% c('aKc1','aKc2','aKc3')) |>
  mutate(v3.type = ifelse(FineClust == 'aKc1', 'V3-lo KC', 'V3-hi KC'))

aobj.kc.v3 <- aobj.kc.v3 |>
  mutate(v3Clust = case_match(FineClust,
                                'aKc1' ~ 'V3-lo bKC',
                                'aKc2' ~ 'V3-lo sKC',
                                .default = 'V3-hi sKC') |>
           fct_relevel('V3-lo bKC','V3-lo sKC'))

aobj.kc.v3 |>
  CoveragePlot(region = c("HMGCS1"), peaks = F,
               extend.downstream = 2000, group.by = 'FineClust') &
  scale_fill_brewer(palette = 'Blues')  &
  theme_jpub &
  NoLegend()

publish_pdf('mission/FPP/ober23/HMGCS1.coverage.scATAC.pdf', width = 40)

aobj.kc.v3 |>
  CoveragePlot(region = c("HMGCR"), peaks = F,
               extend.upstream = 2000, group.by = 'v3Clust') &
  scale_fill_brewer(palette = 'Blues')  &
  theme_jpub &
  NoLegend()

publish_pdf('mission/FPP/ober23/HMGCR.coverage.scATAC.pdf', width = 60)

aobj.kc.v3 |>
  CoveragePlot(region = c("MVK"), peaks = F,
               extend.upstream = 2000, group.by = 'FineClust') &
  scale_fill_brewer(palette = 'Blues')  &
  theme_jpub &
  NoLegend()

publish_pdf('mission/FPP/ober23/MVK.coverage.scATAC.pdf', width = 60)

aobj.kc.v3 |>
  CoveragePlot(region = c("TRPV3"), peaks = F,
               extend.upstream = 2000, group.by = 'FineClust') &
  scale_fill_brewer(palette = 'Blues')  &
  theme_jpub &
  NoLegend()

publish_pdf('TRPV3.coverage.scATAC.pdf', width = 60)

cov.s1 <- aobj.kc.v3 |>
  CoveragePlot(region = c("HMGCS1"), peaks = F, annotation = F,
               extend.downstream = 2000, group.by = 'FineClust') +
  scale_fill_brewer(palette = 'Blues')

ann.s1 <- aobj.kc.v3 |>
  AnnotationPlot(region = 'HMGCS1', extend.downstream = 2000)

til.s1 <- aobj.kc.v3 |>
  TilePlot(region = 'HMGCS1', extend.downstream = 2000)

exp.s1 <- aobj.kc.v3 |>
  ExpressionPlot('HMGCS1', assay = 'RNA', group.by = 'FineClust') +
  scale_fill_brewer(palette = 'Blues')

CombineTracks(list(cov.s1, ann.s1),
              heights = c(5,1), widths = c(5,1)) &
  theme_jpub &
  NoLegend()

publish_pdf('mission/FPP/ober23/HMGCS1.coverage.scATAC.pdf', width = 60)

aobj.kc.v3 |>
  FindMarkers(features = 'HMGCS1', assay = 'RNA',
              group.by = 'FineClust', ident.1 = 'aKc2', ident.2 = 'aKc1')

peak.3v1 <- aobj.kc.v3 |>
  FindMarkers(assay = 'ATAC', group.by = 'FineClust',
              ident.1 = 'aKc3', ident.2 = c('aKc1','aKc2'))

peak.3v1 |>
  as_tibble(rownames = 'region') |>
  separate_wider_delim(region,delim = '-', names = c('chr','start','end')) |>
  mutate(start = as.integer(start), end = as.integer(end)) |>
  filter(chr == 'chr5', start > 4331e4, start < 4333e4)

aobj.kc.v3 |>
  CoveragePlot(region = c("CCL20"), peaks = F,
               extend.upstream = 2000, group.by = 'FineClust')

aobj.kc.v3 |>
  CoveragePlot(region = c("HMGCR"), peaks = F,
               extend.upstream = 2000, group.by = 'FineClust')

aobj.kc.v3 |>
  CoveragePlot(region = c("MVK"), features = 'MVK', tile = T,
               extend.upstream = 2000, group.by = 'FineClust')

aobj.kc.v3 |>
  CoveragePlot(region = c('FDPS'), features = 'FDPS', tile = T,
               group.by = 'v3.type')

aobj.kc <- aobj.kc |>
  mutate(v3.type = case_when(FineClust == 'aKc1' ~'V3-lo KC',
                             FineClust %in% c('aKc2','aKc3') ~'V3-hi KC',
                             .default = 'HF KC'))

aobj.kc |>
  CoveragePlot(region = c("HMGCS1"),
               group.by = 'v3.type',
               features = 'HMGCS1')


aobj.v3hl.fine |>
  ExpressionPlot(
  features = c("HMGCS1"),
  assay = "RNA"
)

## ADRB2 in KC -------
aobj.kc.v3 |>
  CoveragePlot(region = 'ADRB2',
               group.by = 'FineClust')

aobj.kc.v3 |>
  FindMarkers(assay = 'RNA', features = 'ADRB2',
              group.by = 'FineClust', ident.1 = 'aKc3', ident.2 = 'aKc1')

aobj.kc.v3$group |> table()

# rds checkpoint ----------
aobj.kc.v3 |>
  write_rds('mission/FPP/ober23/ober23.v3kc.rds')

aobj.kc.v3 <- read_rds('mission/FPP/ober23/ober23.v3kc.rds')

aobj.kc.v3 <- aobj.kc.v3 |> LinkPeaks(peak.assay = 'ATAC',
                                      expression.assay = 'RNA',
                                      genes.use = c('HMGCS1','HMGCR'))

aobj.kc.v3 |> LinkPlot('HMGCR')

aobj.kc.v3 |> VlnPlot('TRPV3', pt.size = 0)

v3.actv <- aobj.kc.v3 |> get_abundance_sc_wide('TRPV3', assay = 'RNA')

aobj.kc.v3 |>
  left_join(v3.actv) |>
  summarise(avg.exp = logtpm.mean(TRPV3), .by = c(FineClust, dataset)) |>
  tidyplot(FineClust, avg.exp) |>
  add_boxplot()

aobj.kc.v3 |> CoveragePlot('TRPV3', extend.upstream = 2000)

aobj.kc.v3 |> FindMarkers(ident.1 = 'aKc2', features = 'TRPV3', assay = 'RNA',
                          ident.2 = 'aKc1')
